/*********************************************************
          File Name:rb_tree.c
          Author: Abby Cin
          Mail: abbytsing@gmail.com
          Created Time: Fri 08 May 2015 02:40:50 PM CST
**********************************************************/

#include <stdio.h>
#include <stdlib.h>

typedef struct rb{
    struct rb *p, *left, *right;
    enum {red, black} color;
    int key;
    int data;
}rbnode;

typedef struct{
    rbnode *root, *nil;
} rbtree;

void left_rotate(rbtree *T, rbnode *x)
{
    rbnode *y;

    y = x->right;                   /* no matter x->right is NIL or something */
    x->right = y->left;

    if(y->left != T->nil)             /* i.e. x->right->left != NIL */
        y->left->p = x;             /* raise up y's left child */

    y->p = x->p;                    /* raise up y */

    if(x->p == T->nil)
        T->root = y;
    else 
    {
        if(x == x->p->left)
            x->p->left = y;             /* detach x */
        else
            x->p->right = y;
    }

    y->left = x;                    /* drop x */
    x->p = y;                       /* link x to new parents */
}

void right_rotate(rbtree *T, rbnode *x)
{
    rbnode *y;

    y = x->left;
    x->left = y->right;

    if(y->right != T->nil)
        y->right->p = x;

    y->p = x->p;

    if(x->p == T->nil)
        T->root = y;
    else 
    {
        if(x == x->p->right)
            x->p->right = y;
        else
            x->p->left = y;
    }
    y->right = x;
    x->p = y;
}

void rb_insert_fixup(rbtree *T, rbnode *z)
{
    rbnode *y;
    while(z->p->color == red)               /* fix it from buttom to top, raise up z in every loop */
    {
        if(z->p == z->p->p->left)           /* if z->p is left child or right */
        {
            y = z->p->p->right;             /* get uncle of z */
            if(y->color == red)             /* z,z's parents,z's uncle are all red, against 4th law */
            {
                z->p->color     = black;    /* re-color parent to black, protect 4th law */
                y->color        = black;
                z->p->p->color  = red;
                z               = z->p->p;  /* raise up z's pointer */
            }
            else 
            {
                if(z == z->p->right)
                {
                    z = z->p;                   
                    left_rotate(T, z);          /* for z->p->p->p RH, so left_rotate */
                }
                z->p->color     = black;
                z->p->p->color  = red;
                right_rotate(T, z->p->p);   /* for z->p->p->p LH, so right_rotate */
            }
        }
        else                                /* z->p is right child */
        {
            y = z->p->p->left;
            if(y->color == red)
            {
                z->p->color     = black;
                y->color        = black;
                z->p->p->color  = red;
                z               = z->p->p;
            }
            else
            { 
                if(z == z->p->left)
                {
                    z               = z->p;
                    right_rotate(T, z);
                }
                z->p->color     = black;
                z->p->p->color  = red;
                left_rotate(T, z->p->p);
            }
        }
    }
    T->root->color = black;             /* protect 2nd law */
}

void rb_insert(rbtree *T, rbnode *z)
{
    rbnode *y, *x;

    y = T->nil;
    x = T->root;

    while(x != T->nil)                    /* from top to buttom, find a place for instertion */
    {
        y = x;
        if(z->key < x->key)
            x = x->left;
        else
            x = x->right;
    }

    z->p = y;

    if(y == T->nil)
        T->root = z;
    else if(z->key < y->key)
        y->left = z;
    else
        y->right = z;

    z->left = T->nil;
    z->right = T->nil;
    z->color = red;

    rb_insert_fixup(T, z);
}
static rbnode *find_min(rbtree *tree, rbnode *root)
{
  // descend left subtree
    while(root->left != tree->nil)
        root = root->left;
    return root;
}

/*
rbnode *rb_successor(rbtree *T, rbnode *z)
{
    rbnode *tmp = T->root;
    rbnode *y = z;
    if((y = z->right) != T->nil)
    {
        while(y->left != T->nil)
            y = y->left;
        return y;
    }
    else
    {
        y = z->p;
        while(z == y->right)
        {
            z = y;
            y = y->p;
        }
        if(y == T->root) return T->nil;
        return y;
    }
}
*/

static rbnode *rb_successor(rbtree *tree, rbnode *n)
{
  // if we have a right subtree, find minimum in that
    if(n->right != tree->nil)
        n = find_min(tree, n->right);
    else
    {
        rbnode *p = n->p;
        while(p != tree->nil && n == p->right)
        {
            n = p;
            p = n->p;
        }
        n = p;
    } 

    return n;
}

void rb_transplant(rbtree *T, rbnode *u, rbnode *v)
{
    if(u->p == T->nil)
        T->root = v;
    else if(u == u->p->left)
        u->p->left = v;
    else
        u->p->right = v;
    
    v->p = u->p;                /* detach u from tree */
}

void rb_delete_fixup(rbtree *T, rbnode *x)
{
    rbnode *w;

    while(x != T->root && x->color == black)    /* in the satuation, can NOT simply free x */
    {
        if(x == x->p->left)
        {
            w = x->p->right;                    /* get sibling */
            if(w->color == red)
            {
                w->color    = black;
                x->p->color = red;
                left_rotate(T, x->p);
                 w = x->p->right;
            }
            if(w->left->color == black && w->right->color == black) /* w is black */
            {
                w->color = red;
                x        = x->p;
            }
            else 
            {
                if(w->right->color ==  black)                        /* w is black */
                {
                    w->left->color      = black;
                    w->color            = red;
                    right_rotate(T, w);
                    w                   = x->p->right;
                }
                w->color            = x->p->color;
                x->p->color         = black;
                w->right->color     = black;
                left_rotate(T, x->p);
                x                   = T->root;
             }
        }
        else
        {
            w = x->p->left;
            if(w->color == red)
            {
                w->color        = black;
                x->p->color     = red;
                right_rotate(T, x->p);
                w = x->p->left;
            }
            if(w->left->color == black && w->right->color == black)
            {
                w->color        = red;
                x               = x->p;
            }
            else 
            {
                if(w->left->color ==  black)
                {
                    w->right->color     = black;
                    w->color            = red;
                    left_rotate(T, w);
                    w                   = x->p->left;
                }
                w->color            = x->p->color;
                x->p->color         = black;
                w->right->color     = black;
                right_rotate(T, x->p);
                x                   = T->root;
            }
        }
    }

    x->color = black;
}

void rb_delete(rbtree *T, rbnode *z)
{
    rbnode *y, *x; 

    y = z;
    enum {red, black} ori_color;
    ori_color = y->color;

    if(z->left == T->nil)
    {
        x = z->right; 
        rb_transplant(T, z, z->right);
    }
    else if(z->right == T->nil)
    {
        x = z->left; 
        rb_transplant(T, z, z->left);
    }
    else
    {
        y = rb_successor(T, z) ;
        ori_color = y->color;
        x = y->right;
        if(y->p == z)
            x->p = y;
        else
        {
            rb_transplant(T, y, y-> right);
            y->right = z->right;
            y->right->p = y;
        }
        rb_transplant(T, z, y);
        y->left = z->left;
        y->left->p = y;
        y->color = z->color;
    }

    free(z);

    if(ori_color == black)
        rb_delete_fixup(T, x);
}

void inorder_walk(rbtree *T, rbnode *x)
{
    if(x != T->nil)
    {
        inorder_walk(T, x->left);
        printf("key = %d\t\t%-s", x->key, x->color == red ? "red" : "black");
        if(x->p == T->nil)
            printf("\e[34m\troot\e[0m\n");
        else
            printf("\n");
        inorder_walk(T, x->right);
    }
}

void rb_destroy(rbtree *T, rbnode *root)
{
    rbnode *tmp = root;
    if(tmp != T->nil)
    {
        if(tmp->left != T->nil)
            rb_destroy(T, tmp->left);
        if(tmp->right != T->nil)
            rb_destroy(T, tmp->right);
        printf("free %d\t\t%-s\n", tmp->key, tmp->color == red ? "red" : "black");
        free(tmp);
    }
    free(T->nil);
}

void rb_build_tree(rbtree *T, int size)
{
    rbnode *tmp;
    int i;

    for(i=0; i<size; i++)
    {
        tmp = (rbnode *)malloc(sizeof(rbnode));
        if(tmp == NULL)
        {
            perror("malloc");
            exit(1);
        }
        printf("enter a key:");
        scanf("%d", &tmp->key);
        rb_insert(T, tmp);
    }
}

int main(int argc, char **argv)
{
    if(argc != 2)
    {
        fprintf(stderr, "%s nodes\n", argv[0]);
        exit(1);
    }
    rbtree *T;
    T = (rbtree *)malloc(sizeof(rbtree));
    T->nil = (rbnode *)malloc(sizeof(rbnode));
    T->nil->color = black;
    T->root = T->nil;

    rb_build_tree(T, atoi(argv[1]));
    inorder_walk(T, T->root);

    rbnode *tmp = NULL;
    tmp = (rbnode *)malloc(sizeof(rbnode));
    printf("enter a key to insert: ");
    scanf("%d", &tmp->key);
    rb_insert(T, tmp);
    inorder_walk(T, T->root);

    printf("after deleting: %d\n", T->root->right->key);
    rb_delete(T, T->root->right);
    inorder_walk(T, T->root);

    rb_destroy(T, T->root);
    free(T);
}
